# encoding:utf-8

"""
wechat channel
"""

import os
import requests
import sqlite3
import io
import time
import json
import datetime
from channel.chat_channel import ChatChannel
from channel.wechat.wechat_message import *
from common.singleton import singleton
from common.log import logger
from lib import itchat
from lib.itchat.content import *
from bridge.reply import *
from bridge.context import *
from concurrent.futures import ThreadPoolExecutor
from config import conf
from common.time_check import time_checker
from common.expired_dict import ExpiredDict
from plugins import *
thread_pool = ThreadPoolExecutor(max_workers=8)

def thread_pool_callback(worker):
    worker_exception = worker.exception()
    if worker_exception:
        logger.exception("Worker return exception: {}".format(worker_exception))

# 注册信息
@itchat.msg_register(MAP, isGroupChat=True)
def handler_single_msg(msg):
    logger.info("map msg", msg)
    return None
@itchat.msg_register(CARD, isGroupChat=True)
def handler_single_msg(msg):
    logger.info("CARD msg", msg)
    return None
@itchat.msg_register(SHARING, isGroupChat=True)
def handler_single_msg(msg):
    logger.info("SHARING msg", msg)
    return None
@itchat.msg_register(PICTURE, isGroupChat=True)
def handler_single_msg(msg):
    logger.info("PICTURE msg", msg)
    return None
@itchat.msg_register(ATTACHMENT, isGroupChat=True)
def handler_single_msg(msg):
    logger.info("ATTACHMENT msg", msg)
    return None
@itchat.msg_register(VIDEO, isGroupChat=True)
def handler_single_msg(msg):
    logger.info("VIDEO msg", msg)
    return None
@itchat.msg_register(FRIENDS, isGroupChat=True)
def handler_single_msg(msg):
    logger.info("FRIENDS msg", msg)
    return None
@itchat.msg_register(SYSTEM, isGroupChat=True)
def handler_single_msg(msg):
    logger.info("SYSTEM msg", msg)
    return None

# 邀请进群事件回掉在这里
@itchat.msg_register(NOTE, isGroupChat=True)
def handler_single_msg(msg):
    logger.info("NOTE msg", msg)
    WechatChannel().handle_invite(WeChatMessage(msg, True))
    return None

@itchat.msg_register(TEXT)
def handler_single_msg(msg):
    WechatChannel().handle_text(WeChatMessage(msg))
    return None

@itchat.msg_register(TEXT, isGroupChat=True)
def handler_group_msg(msg):
    WechatChannel().handle_group(WeChatMessage(msg,True))
    return None

@itchat.msg_register(VOICE)
def handler_single_voice(msg):
    WechatChannel().handle_voice(WeChatMessage(msg))
    return None
    
@itchat.msg_register(VOICE, isGroupChat=True)
def handler_group_voice(msg):
    WechatChannel().handle_group_voice(WeChatMessage(msg,True))
    return None

def _check(func):
    def wrapper(self, cmsg: ChatMessage):
        msgId = cmsg.msg_id
        if msgId in self.receivedMsgs:
            logger.info("Wechat message {} already received, ignore".format(msgId))
            return
        self.receivedMsgs[msgId] = cmsg
        create_time = cmsg.create_time            # 消息时间戳
        if conf().get('hot_reload') == True and int(create_time) < int(time.time()) - 60:  # 跳过1分钟前的历史消息
            logger.debug("[WX]history message {} skipped".format(msgId))
            return
        return func(self, cmsg)
    return wrapper

@singleton
class WechatChannel(ChatChannel):
    def __init__(self):
        super().__init__()
        self.receivedMsgs = ExpiredDict(60*60*24)

    def startup(self):

        itchat.instance.receivingRetryCount = 600  # 修改断线超时时间
        # login by scan QRCode
        hotReload = conf().get('hot_reload', True)
        try:
            itchat.auto_login(enableCmdQR=2, hotReload=hotReload)
        except Exception as e:
            if hotReload:
                logger.error("Hot reload failed, try to login without hot reload")
                itchat.logout()
                os.remove("itchat.pkl")
                itchat.auto_login(enableCmdQR=2, hotReload=hotReload)
            else:
                raise e
        self.user_id = itchat.instance.storageClass.userName
        self.name = itchat.instance.storageClass.nickName
        logger.info("Wechat login success, user_id: {}, nickname: {}".format(self.user_id, self.name))
        # start message listener
        itchat.run()

    # handle_* 系列函数处理收到的消息后构造Context，然后传入_handle函数中处理Context和发送回复
    # Context包含了消息的所有信息，包括以下属性
    #   type 消息类型, 包括TEXT、VOICE、IMAGE_CREATE
    #   content 消息内容，如果是TEXT类型，content就是文本内容，如果是VOICE类型，content就是语音文件名，如果是IMAGE_CREATE类型，content就是图片生成命令
    #   kwargs 附加参数字典，包含以下的key：
    #        session_id: 会话id
    #        isgroup: 是否是群聊
    #        receiver: 需要回复的对象
    #        msg: ChatMessage消息对象
    #        origin_ctype: 原始消息类型，语音转文字后，私聊时如果匹配前缀失败，会根据初始消息是否是语音来放宽触发规则
    #        desire_rtype: 希望回复类型，默认是文本回复，设置为ReplyType.VOICE是语音回复

    @time_checker
    @_check
    def handle_invite(self, cmsg : ChatMessage):
        config=conf()
        if "群聊" not in cmsg.content:
            return
        # 进群欢迎文案
        context = Context(ContextType.TEXT, "邀请进群提示")
        # 邀请人
        yaoqingren_username = ""
        # 被邀请人
        beiyaoqingren_username = ""
        invite_welcome_tips=config.get('invite_welcome_tips', "欢迎新人进群，我是本群智能机器人bot，你可以随时向我提问")
        invite_reward_num=int(config.get('invite_reward_num', 10))
        # 回复兜底信息
        if "邀请" in cmsg.content:
            reply = Reply(ReplyType.INFO, '@' + cmsg.content.replace("加入了群聊", "").split("邀请")[1].replace('"', "") + " " + invite_welcome_tips)
            context["receiver"] = cmsg.other_user_id
            yaoqingren_username = cmsg.content.replace("加入了群聊", "").split("邀请")[0].replace('"', "")
            beiyaoqingren_username = cmsg.content.replace("加入了群聊", "").split("邀请")[1].replace('"', "")
            # reply的发送步骤
            self.send(reply, context)
        elif "通过扫描" in cmsg.content:
            reply = Reply(ReplyType.INFO, '@' + cmsg.content.replace("分享的二维码加入群聊", "").split("通过扫描")[0].replace('"', "") + " " + invite_welcome_tips)
            context["receiver"] = cmsg.other_user_id
            yaoqingren_username = cmsg.content.replace("分享的二维码加入群聊", "").split("通过扫描")[1].replace('"', "")
            beiyaoqingren_username = cmsg.content.replace("分享的二维码加入群聊", "").split("通过扫描")[0].replace('"', "")
            # reply的发送步骤
            self.send(reply, context)
        # 次数限制逻辑，仅限群聊
        conn = sqlite3.connect("/home/ubuntu/apps/chatgpt-admin/db.sqlite3")
        cursor = conn.cursor()
        # todo 查询历史是否邀请过，根据 reward_reason 字段
        logger.info(str(cmsg.from_user_id))

        # 查询用户邀请了几个人进群
        cursor.execute("SELECT * FROM tb_reward_num where user_id='{}'".format(yaoqingren_username + cmsg.from_user_id.replace("@", "")))
        invite_info= {}
        invite_info[cmsg.from_user_id.replace("@", "")] = cursor.fetchall()

        if len(invite_info[cmsg.from_user_id.replace("@", "")]) == 0:
            # 插入进群奖励次数
            cursor.execute(
                "INSERT INTO tb_reward_num ('user_id', 'reward_num', 'reward_reason', 'create_time') VALUES ('{}', {}, '{}', '{}')".format(
                    yaoqingren_username + cmsg.from_user_id.replace("@", ""),
                    invite_reward_num,
                    "invite_reward: from_user_id + to_user_id",
                    datetime.datetime.now()
                ))
            conn.commit()
            # 更新今天的使用次数
            cursor.execute("SELECT * FROM tb_chat_usage_recorder where receiver_id='{}' and session_id='{}' order by create_time desc".format(cmsg.from_user_id.replace("@", ""), yaoqingren_username))
            history_recorders=cursor.fetchall()
            new_total_num = invite_reward_num
            if len(history_recorders) > 0:
                new_total_num=history_recorders[0][3] + new_total_num
            cursor.execute("UPDATE tb_chat_usage_recorder SET 'total_num' = {} WHERE receiver_id = '{}' and session_id='{}'".format(str(new_total_num), cmsg.from_user_id.replace("@", ""), yaoqingren_username))
            conn.commit()

        elif invite_info[cmsg.from_user_id.replace("@", "")][0][2] == int(2 * invite_reward_num):
            # 在原先的基础上加 10 次
            cursor.execute("SELECT * FROM tb_reward_num where user_id='{}'".format(yaoqingren_username + cmsg.from_user_id.replace("@", "")))
            reward_num = cursor.fetchall()[0][2]
            cursor.execute(
                "UPDATE tb_reward_num SET 'reward_num' = {} WHERE user_id = '{}'".format(
                    str(int(reward_num) + invite_reward_num),
                    yaoqingren_username + cmsg.from_user_id.replace("@", "")
                )
            )
            conn.commit()

            # 自定义提醒
            context = Context(ContextType.TEXT, "邀请进群提示")
            invite_three_tips = config.get('invite_three_tips', "累计邀请好友进群3人，获得本群同款机器人一个，可截图找群主领取~")
            # 回复兜底信息
            if "邀请" in cmsg.content:
                reply = Reply(ReplyType.INFO, '@' + cmsg.content.replace("加入了群聊", "").split("邀请")[0].replace('"', "") + " " + invite_three_tips)
                context["receiver"] = cmsg.other_user_id
                # reply的发送步骤
                self.send(reply, context)
            elif "通过扫描" in cmsg.content:
                reply = Reply(ReplyType.INFO, '@' + cmsg.content.replace("分享的二维码加入群聊", "").split("通过扫描")[1].replace('"', "") + " " + invite_three_tips)
                context["receiver"] = cmsg.other_user_id
                # reply的发送步骤
                self.send(reply, context)

            # 更新今天的使用次数
            cursor.execute("SELECT * FROM tb_chat_usage_recorder where receiver_id='{}' and session_id='{}' order by create_time desc".format(cmsg.from_user_id.replace("@", ""), yaoqingren_username))
            history_recorders=cursor.fetchall()
            new_total_num = invite_reward_num
            if len(history_recorders) > 0:
                new_total_num=history_recorders[0][3] + new_total_num
            cursor.execute("UPDATE tb_chat_usage_recorder SET 'total_num' = {} WHERE receiver_id = '{}' and session_id='{}'".format(str(new_total_num), cmsg.from_user_id.replace("@", ""), yaoqingren_username))
            conn.commit()

        elif invite_info[cmsg.from_user_id.replace("@", "")][0][2] == int(4 * invite_reward_num):
            # 在原先的基础上加 10 次
            cursor.execute("SELECT * FROM tb_reward_num where user_id='{}'".format(yaoqingren_username + cmsg.from_user_id.replace("@", "")))
            reward_num = cursor.fetchall()[0][2]
            cursor.execute(
                "UPDATE tb_reward_num SET 'reward_num' = {} WHERE user_id = '{}'".format(
                    str(int(reward_num) + invite_reward_num),
                    yaoqingren_username + cmsg.from_user_id.replace("@", "")
                )
            )
            conn.commit()
            # 自定义提醒
            context = Context(ContextType.TEXT, "邀请进群提示")
            invite_five_tips = config.get('invite_five_tips', "累计邀请好友进群5人，奖励私人机器人一个（无敏感词和次数限制），可截图找群主领取~")
            # 回复兜底信息
            if "邀请" in cmsg.content:
                reply = Reply(ReplyType.INFO, '@' + cmsg.content.replace("加入了群聊", "").split("邀请")[0].replace('"', "") + " " + invite_five_tips)
                context["receiver"] = cmsg.other_user_id
                # reply的发送步骤
                self.send(reply, context)
            elif "通过扫描" in cmsg.content:
                reply = Reply(ReplyType.INFO, '@' + cmsg.content.replace("分享的二维码加入群聊", "").split("通过扫描")[1].replace('"', "") + " " + invite_five_tips)
                context["receiver"] = cmsg.other_user_id
                # reply的发送步骤
                self.send(reply, context)
            # 更新今天的使用次数
            cursor.execute("SELECT * FROM tb_chat_usage_recorder where receiver_id='{}' and session_id='{}' order by create_time desc".format(cmsg.from_user_id.replace("@", ""), yaoqingren_username))
            history_recorders=cursor.fetchall()
            new_total_num = invite_reward_num
            if len(history_recorders) > 0:
                new_total_num=history_recorders[0][3] + new_total_num
            cursor.execute("UPDATE tb_chat_usage_recorder SET 'total_num' = {} WHERE receiver_id = '{}' and session_id='{}'".format(str(new_total_num), cmsg.from_user_id.replace("@", ""), yaoqingren_username))
            conn.commit()

        else:
            # 在原先的基础上加 10 次
            cursor.execute("SELECT * FROM tb_reward_num where user_id='{}'".format(yaoqingren_username + cmsg.from_user_id.replace("@", "")))
            reward_num = cursor.fetchall()[0][2]
            cursor.execute(
                "UPDATE tb_reward_num SET 'reward_num' = {} WHERE user_id = '{}'".format(
                    str(int(reward_num) + invite_reward_num),
                    yaoqingren_username + cmsg.from_user_id.replace("@", "")
                )
            )
            conn.commit()
            # 更新今天的使用次数
            cursor.execute("SELECT * FROM tb_chat_usage_recorder where receiver_id='{}' and session_id='{}' order by create_time desc".format(cmsg.from_user_id.replace("@", ""), yaoqingren_username))
            history_recorders=cursor.fetchall()
            new_total_num = invite_reward_num
            if len(history_recorders) > 0:
                new_total_num=history_recorders[0][3] + new_total_num
            cursor.execute("UPDATE tb_chat_usage_recorder SET 'total_num' = {} WHERE receiver_id = '{}' and session_id='{}'".format(str(new_total_num), cmsg.from_user_id.replace("@", ""), yaoqingren_username))
            conn.commit()

        cursor.close()
        conn.close()


    @time_checker
    @_check
    def handle_voice(self, cmsg : ChatMessage):
        if conf().get('speech_recognition') != True:
            return
        logger.debug("[WX]receive voice msg: {}".format(cmsg.content))
        context = self._compose_context(ContextType.VOICE, cmsg.content, isgroup=False, msg=cmsg)
        if context:
            thread_pool.submit(self._handle, context).add_done_callback(thread_pool_callback)

    @time_checker
    @_check
    def handle_text(self, cmsg : ChatMessage):
        logger.debug("[WX]receive text msg: {}, cmsg={}".format(json.dumps(cmsg._rawmsg, ensure_ascii=False), cmsg))
        context = self._compose_context(ContextType.TEXT, cmsg.content, isgroup=False, msg=cmsg)
        if context:
            thread_pool.submit(self._handle, context).add_done_callback(thread_pool_callback)

    @time_checker
    @_check
    def handle_group(self, cmsg : ChatMessage):
        logger.debug("[WX]receive group msg: {}, cmsg={}".format(json.dumps(cmsg._rawmsg, ensure_ascii=False), cmsg))
        context = self._compose_context(ContextType.TEXT, cmsg.content, isgroup=True, msg=cmsg)
        if context:
            thread_pool.submit(self._handle, context).add_done_callback(thread_pool_callback)
    
    @time_checker
    @_check
    def handle_group_voice(self, cmsg : ChatMessage):
        if conf().get('group_speech_recognition', False) != True:
            return
        logger.debug("[WX]receive voice for group msg: {}".format(cmsg.content))
        context = self._compose_context(ContextType.VOICE, cmsg.content, isgroup=True, msg=cmsg)
        if context:
            thread_pool.submit(self._handle, context).add_done_callback(thread_pool_callback)
    
    # 统一的发送函数，每个Channel自行实现，根据reply的type字段发送不同类型的消息
    def send(self, reply: Reply, context: Context):
        receiver = context["receiver"]
        if reply.type == ReplyType.TEXT:
            itchat.send(reply.content, toUserName=receiver)
            logger.info('[WX] sendMsg={}, receiver={}'.format(reply, receiver))
        elif reply.type == ReplyType.ERROR or reply.type == ReplyType.INFO:
            itchat.send(reply.content, toUserName=receiver)
            logger.info('[WX] sendMsg={}, receiver={}'.format(reply, receiver))
        elif reply.type == ReplyType.VOICE:
            itchat.send_file(reply.content, toUserName=receiver)
            logger.info('[WX] sendFile={}, receiver={}'.format(reply.content, receiver))
        elif reply.type == ReplyType.IMAGE_URL: # 从网络下载图片
            img_url = reply.content
            pic_res = requests.get(img_url, stream=True)
            image_storage = io.BytesIO()
            for block in pic_res.iter_content(1024):
                image_storage.write(block)
            image_storage.seek(0)
            itchat.send_image(image_storage, toUserName=receiver)
            logger.info('[WX] sendImage url={}, receiver={}'.format(img_url,receiver))
        elif reply.type == ReplyType.IMAGE: # 从文件读取图片
            image_storage = reply.content
            image_storage.seek(0)
            itchat.send_image(image_storage, toUserName=receiver)
            logger.info('[WX] sendImage, receiver={}'.format(receiver))


